from os import name
import torch
from torch import nn
from torch.optim import optimizer

def regularization(named_parameter, state_dict, device = 'cuda'):
    total_value = 0
    for name, param in named_parameter:
        state_dict[name] = state_dict[name].to(device)
        param = param.to(device)
        total_value += torch.norm(state_dict[name] - param, p = 2)
    return total_value


class TrainModel:
    def __init__(self, model) -> None:
        self.model = model

    def get_model_params(self):
        # return self.model.cpu().state_dict(),要不要转移到cpu上去呢？
        return self.model.cpu().state_dict()
    
    def set_model_params(self, model_params):
        self.model.load_state_dict(model_params)

    def train(self, train_data, epoch, lr, device = 'cuda'):
        # model = copy.deepcopy(self.model)  千万不要用这一句，找了两天的bug！！
        #如果用这一句，self.model不变，因此get_model_params得到的永远是同一个params
        model = self.model
        model.to(device)
        model.train()

        loss_fn = nn.CrossEntropyLoss().to(device)
        optimizer = torch.optim.SGD(model.parameters(), lr = lr)
        epoch = epoch

        for _ in range(epoch):
            for idx, (x, y) in enumerate(train_data):
                x, y = x.to(device), y.to(device)
                #这里的数据x，y在load_partition_data的时候就已经是tensor了 
                output = model(x)
                loss = loss_fn(output, y.long())
                optimizer.zero_grad()
                loss.backward()
                # to avoid nan loss
                #torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()


    def test(self, test_data, device = 'cuda'):
        # model = copy.deepcopy(self.model)
        model = self.model
        model.to(device)
        model.eval()

        metrics = {
            'test_correct' : 0,
            'test_loss' : 0,
            'test_total' : 0
        }

        loss_fn = nn.CrossEntropyLoss().to(device)
        with torch.no_grad():
            for idx, (x, y) in enumerate(test_data):
                x = x.to(device)
                y = y.to(device)
                pred = model(x)
                loss = loss_fn(pred, y.long())

                _, max_location = torch.max(pred, -1)
                correct_num = max_location.eq(y).sum()
                
                metrics['test_correct'] += correct_num.item()
                metrics['test_loss'] += loss.item() * y.size(0)
                metrics['test_total'] += y.size(0)
        return metrics

    def get_y(self, data, device = 'cuda'):
        model = self.model
        model.to(device)
        model.eval()
        output = []
        with torch.no_grad():
            for idx, (x, y) in enumerate(data):
                x = x.to(device)
                # y = y.to(device)
                pred = model(x)
                output.append(pred)
        
        return output


class per_TrainModel(TrainModel):
    def __init__(self, model) -> None:
        super(per_TrainModel, self).__init__(model)

    def set_model_params(self, model_params):
        self.model.load_state_dict(model_params, strict = False)

    def get_model_params(self,per_layer, model_layer = 6):
        model_state_dict = self.model.state_dict()
        per_layer_num = per_layer * 2
        model_layer_num = model_layer * 2
        return_layer_num = model_layer_num - per_layer_num
        key_list = []
        for key in model_state_dict:
            key_list.append(key)
        for key in key_list[return_layer_num :]:
            model_state_dict.pop(key)

        return model_state_dict

    def Cper_train(self, train_data, Cper_epoch, Cper_lr, state_dict, device = 'cuda'):
        model = self.model
        model.to(device)
        model.train()

        loss_fn = nn.CrossEntropyLoss().to(device)
        optimizer = torch.optim.SGD(model.parameters(), lr = Cper_lr)

        for _ in range(Cper_epoch):
            for idx, (x, y) in enumerate(train_data):
                x, y = x.to(device), y.to(device)
                #这里的数据x，y在load_partition_data的时候就已经是tensor了 
                output = model(x)

                loss_regularization = regularization(named_parameter = model.named_parameters(), state_dict = state_dict)
                loss = loss_fn(output, y.long()) + loss_regularization
                optimizer.zero_grad()
                loss.backward()
                # to avoid nan loss
                #torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
    
    def Cper_test(self, test_data, device = 'cuda'):
        pass